# img_paths = ['ultralytics/assets/bus.jpg','ultralytics/assets/zidane.jpg']
import requests
import json

from util_for_argparse import get_cuda_port_args

args = get_cuda_port_args()

def test_get_pose_by_yolo_batch(server_url=f'http://127.0.0.1:{args.port}'):
    # 测试通过url列表方式
    img_paths = ['ultralytics/assets/bus.jpg', 'ultralytics/assets/zidane.jpg']
    url = f'{server_url}/api/get_pose_by_yolo_batch'
    payload = {
        'img_paths': img_paths
    }
    headers = {'Content-Type': 'application/json'}
    resp = requests.post(url, data=json.dumps(payload), headers=headers)
    print('--- Batch URL list result ---')
    print('status:', resp.status_code)
    # print(resp.json())
    print()

    # 测试通过文件上传方式
    files = [
        ('images', ('bus.jpg', open('ultralytics/assets/bus.jpg', 'rb'), 'image/jpeg')),
        ('images', ('zidane.jpg', open('ultralytics/assets/zidane.jpg', 'rb'), 'image/jpeg'))
    ]
    resp2 = requests.post(url, files=files)
    print('--- Batch File Upload result ---')
    print('status:', resp2.status_code)
    # print(resp2.json())

    return resp2.json()

def test_draw_pose(server_url=f'http://127.0.0.1:{args.port}',kps=None):
    # 这里假设有关键点数据kps（形如[[x,y,score],...]），可以用上面接口返回的keypoints
    # 方便测试：读取一张图片并随机dummy keypoints
    import numpy as np
    import cv2

    img_path = 'ultralytics/assets/bus.jpg'
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    # 生成一个人的dummy keypoints
    if kps is None:
        kp_num = 17
        np.random.seed(42)
        kps = np.zeros((kp_num, 3))
        kps[:, 0] = np.random.rand(kp_num) * w
        kps[:, 1] = np.random.rand(kp_num) * h
        kps[:, 2] = np.random.rand(kp_num) * 0.6 + 0.4  # score between 0.4~1

    url = f'{server_url}/api/draw_pose'
    files = {
        'image': open(img_path, 'rb'),
    }
    # 注意kps使用字符串形式提交
    data = {
        'kps': str(kps.tolist()),
        'threshold': 0.3,
        'radius': 5
    }
    resp = requests.post(url, files=files, data=data)
    print('--- Draw Pose result ---')
    print('status:', resp.status_code)
    if resp.status_code == 200:
        # 保存结果图片
        with open('test_pose_result.jpg', 'wb') as f:
            f.write(resp.content)
        print('Saved test_pose_result.jpg')
    else:
        print(resp.text)

if __name__ == '__main__':
    # 运行前请确保flask服务已开启
    results = test_get_pose_by_yolo_batch()

    import numpy as np
    kps = np.array(results['keypoints'][0])

    test_draw_pose(kps=kps)
